# Description: 2AI Message
from langchain_ollama import ChatOllama
from langchain_core.prompts import ChatPromptTemplate

from langchain_core.output_parsers import StrOutputParser
messages = [
    (
        "system",
        "You are a helpful assistant that translates {input_language} to {output_language}."
    ),
    (
        "human",
        "{input}"
    )
]

# 应用上面的 messages
prompt = ChatPromptTemplate.from_messages(messages)
parser = StrOutputParser()


# 创建一个 ChatOllama 实例
llm = ChatOllama(
    model="llama3",
    temperature=0.5,
)

chain = prompt | llm | parser

ai_msg = response = chain.invoke( {
        "input_language": "English",
        "output_language": "Chinese",
        "input": "Hello, how are you?",
    })

print(ai_msg)
